#include "temporal_graph.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <QHash>
#include <QSet>
#include <QVector>

namespace application {
namespace pathfinding {

namespace {

struct StepKey {
    int agvNumber{0};
    int routeIndex{0};

    bool operator==(const StepKey& other) const noexcept {
        return agvNumber == other.agvNumber && routeIndex == other.routeIndex;
    }
};

struct StepKeyHash {
    std::size_t operator()(const StepKey& key) const noexcept {
        std::size_t seed = std::hash<int>{}(key.agvNumber);
        seed ^= (std::hash<int>{}(key.routeIndex) << 1);
        return seed;
    }
};

struct MoveInfo {
    StepKey fromKey;
    StepKey toKey;
    QString fromStation;
    QString toStation;
    int leaveMarkerId{-1};
    int arriveMarkerId{-1};
    double distance{0.0};
};

struct ArrivalInfo {
    int agv{0};
    int routeIndex{0};
    int time{0};
};

struct WaitRelation {
    StepKey awaited;
    StepKey waiter;
    QString station;
    std::size_t constraintIndex{0};
    int awaitedEventId{0};
    int waiterEventId{0};
};

double effectiveSpeed(double configuredSpeed) {
    return configuredSpeed > 0.0 ? configuredSpeed : 1.0;
}

double clampSafetyDistance(double configuredDistance, double moveDistance) {
    if (configuredDistance <= 0.0) {
        return 0.0;
    }
    if (moveDistance <= 0.0) {
        return 0.0;
    }
    const double upper = moveDistance * 0.5;
    return std::min(configuredDistance, upper);
}

void normalizeRouteEntries(const QVector<QString>& originalRoute,
                           const QVector<int>& originalTimes,
                           QVector<QString>& normalizedRoute,
                           QVector<int>& normalizedTimes) {
    normalizedRoute.clear();
    normalizedTimes.clear();
    normalizedRoute.reserve(originalRoute.size());
    normalizedTimes.reserve(originalRoute.size());

    for (int idx = 0; idx < originalRoute.size(); ++idx) {
        const QString& station = originalRoute[idx];
        const int timestamp = (idx < originalTimes.size()) ? originalTimes[idx] : idx;

        if (!normalizedRoute.isEmpty()) {
            const QString& lastStation = normalizedRoute.back();

            if (station == lastStation) {
                normalizedTimes.back() = std::max(normalizedTimes.back(), timestamp);
                continue;
            }

            if (normalizedRoute.size() >= 2) {
                const QString& previousStation = normalizedRoute[normalizedRoute.size() - 2];
                const bool isImmediateReversal =
                    (station == previousStation && lastStation != station);
                const bool upcomingReturnsToLast =
                    (idx + 1 < originalRoute.size() && originalRoute[idx + 1] == lastStation);

                if (isImmediateReversal && upcomingReturnsToLast) {
                    normalizedTimes.back() = std::max(normalizedTimes.back(), timestamp);
                    continue;
                }
            }
        }

        normalizedRoute.append(station);
        normalizedTimes.append(timestamp);
    }
}

}  // namespace

TemporalGraphResult TemporalGraph::build(const std::vector<AgentPlan>& plans,
                                         const TemporalGraphBuildConfig& config) {
    TemporalGraphResult result;
    if (plans.empty()) {
        return result;
    }

    const double speed = effectiveSpeed(config.maxSpeed);

    std::vector<QVector<QString>> normalizedRoutes(plans.size());
    std::vector<QVector<int>> normalizedTimes(plans.size());
    for (size_t planIndex = 0; planIndex < plans.size(); ++planIndex) {
        const auto& plan = plans[planIndex];
        normalizeRouteEntries(plan.route, plan.times,
                              normalizedRoutes[planIndex],
                              normalizedTimes[planIndex]);
    }

    std::vector<TemporalEvent> events;
    std::vector<TemporalConstraint> constraints;
    std::vector<MoveInfo> moves;

    std::unordered_map<StepKey, int, StepKeyHash> stepToEventId;
    std::unordered_map<StepKey, int, StepKeyHash> leaveMarkerByStep;
    std::unordered_map<StepKey, int, StepKeyHash> enterMarkerByStep;
    std::unordered_map<int, std::size_t> leaveMarkerToMoveIndex;

    QHash<QString, QVector<ArrivalInfo>> stationArrivals;
    QHash<QString, QSet<int>> stationUsage;
    QSet<QString> repeatedStations;

    int nextEventId = 0;
    const auto eventIdFor = [&](const StepKey& key) -> int {
        auto it = stepToEventId.find(key);
        return it != stepToEventId.end() ? it->second : -1;
    };
    events.reserve(plans.size() * 4);  // heuristic

    for (size_t planIndex = 0; planIndex < plans.size(); ++planIndex) {
        const auto& plan = plans[planIndex];
        const auto& route = normalizedRoutes[planIndex];
        const auto& routeTimes = normalizedTimes[planIndex];
        QHash<QString, int> localFreq;
        QList<QList<QString>> segments;
        if (route.isEmpty()) {
            result.segmentsByAgv.insert(plan.agvNumber, segments);
            continue;
        }

        for (int idx = 0; idx < route.size(); ++idx) {
            TemporalEvent evt;
            evt.id = nextEventId++;
            evt.agvNumber = plan.agvNumber;
            evt.stationId = route[idx];
            evt.routeIndex = idx;
            evt.type = TemporalEventType::LocationEntry;
            events.push_back(evt);

            StepKey key{plan.agvNumber, idx};
            stepToEventId.emplace(key, evt.id);

            localFreq[evt.stationId] = localFreq.value(evt.stationId, 0) + 1;
            const int timeStamp = (idx < routeTimes.size()) ? routeTimes[idx] : idx;
            stationArrivals[evt.stationId].append({plan.agvNumber, idx, timeStamp});
            stationUsage[evt.stationId].insert(plan.agvNumber);

            if (idx > 0) {
                const QString& prev = route[idx - 1];
                const QString& curr = route[idx];
                if (prev != curr) {
                    QList<QString> seg;
                    seg.append(prev);
                    seg.append(curr);
                    segments.append(seg);
                } else {
                    TemporalConstraint waitConstraint;
                    StepKey fromKey{plan.agvNumber, idx - 1};
                    waitConstraint.fromEventId = eventIdFor(fromKey);
                    waitConstraint.toEventId = evt.id;
                    if (waitConstraint.fromEventId >= 0) {
                        waitConstraint.type = TemporalConstraintType::IntraAgent;
                        waitConstraint.lowerBound = 0.0;
                        constraints.push_back(waitConstraint);
                    }
                }
            }
        }

        for (auto it = localFreq.cbegin(); it != localFreq.cend(); ++it) {
            if (it.value() > 1) {
                repeatedStations.insert(it.key());
            }
        }
        result.segmentsByAgv.insert(plan.agvNumber, segments);
    }

    for (auto it = stationUsage.cbegin(); it != stationUsage.cend(); ++it) {
        if (it.value().size() > 1) {
            result.specialStationWhiteList.insert(it.key(), 1);
        }
    }
    for (const auto& repeated : repeatedStations) {
        result.specialStationWhiteList.insert(repeated, 1);
    }

    std::vector<WaitRelation> waitRelations;

    for (auto it = stationArrivals.begin(); it != stationArrivals.end(); ++it) {
        auto arrivals = it.value();
        std::sort(arrivals.begin(), arrivals.end(),
            [](const ArrivalInfo& lhs, const ArrivalInfo& rhs) {
                if (lhs.time != rhs.time) return lhs.time < rhs.time;
                if (lhs.routeIndex != rhs.routeIndex) return lhs.routeIndex < rhs.routeIndex;
                return lhs.agv < rhs.agv;
            });

        for (int idx = 1; idx < arrivals.size(); ++idx) {
            const ArrivalInfo& awaited = arrivals[idx - 1];
            const ArrivalInfo& waiter = arrivals[idx];
            if (awaited.agv == waiter.agv) {
                continue;
            }

            StepKey awaitedKey{ awaited.agv, awaited.routeIndex };
            StepKey waiterKey{ waiter.agv, waiter.routeIndex };
            const int awaitedEvent = eventIdFor(awaitedKey);
            const int waiterEvent = eventIdFor(waiterKey);
            if (awaitedEvent < 0 || waiterEvent < 0) {
                continue;
            }

            TemporalConstraint interConstraint;
            interConstraint.fromEventId = awaitedEvent;
            interConstraint.toEventId = waiterEvent;
            interConstraint.type = TemporalConstraintType::InterAgent;
            interConstraint.lowerBound = 0.0;
            constraints.push_back(interConstraint);

            WaitRelation relation;
            relation.awaited = awaitedKey;
            relation.waiter = waiterKey;
            relation.station = it.key();
            relation.constraintIndex = constraints.size() - 1;
            relation.awaitedEventId = awaitedEvent;
            relation.waiterEventId = waiterEvent;
            waitRelations.push_back(relation);
        }

        result.specialStationWhiteList.insert(it.key(), 1);
    }

    for (size_t planIndex = 0; planIndex < plans.size(); ++planIndex) {
        const auto& plan = plans[planIndex];
        const auto& route = normalizedRoutes[planIndex];
        if (route.size() < 2) {
            continue;
        }
        for (int idx = 1; idx < route.size(); ++idx) {
            const QString& fromStation = route[idx - 1];
            const QString& toStation = route[idx];
            if (fromStation == toStation) {
                continue;
            }

            StepKey fromKey{plan.agvNumber, idx - 1};
            StepKey toKey{plan.agvNumber, idx};
            const int fromEventId = eventIdFor(fromKey);
            const int toEventId = eventIdFor(toKey);
            if (fromEventId < 0 || toEventId < 0) {
                continue;
            }

            MoveInfo move;
            move.fromKey = fromKey;
            move.toKey = toKey;
            move.fromStation = fromStation;
            move.toStation = toStation;

            double distance = 1.0;
            if (config.distanceProvider) {
                distance = config.distanceProvider(fromStation, toStation);
                if (!(distance > 0.0)) {
                    distance = 1.0;
                }
            }
            move.distance = distance;

            TemporalEvent departMarker;
            departMarker.id = nextEventId++;
            departMarker.agvNumber = plan.agvNumber;
            departMarker.stationId = fromStation;
            departMarker.routeIndex = idx - 1;
            departMarker.type = TemporalEventType::SafetyDepart;
            events.push_back(departMarker);

            TemporalEvent arriveMarker;
            arriveMarker.id = nextEventId++;
            arriveMarker.agvNumber = plan.agvNumber;
            arriveMarker.stationId = toStation;
            arriveMarker.routeIndex = idx;
            arriveMarker.type = TemporalEventType::SafetyArrive;
            events.push_back(arriveMarker);

            move.leaveMarkerId = departMarker.id;
            move.arriveMarkerId = arriveMarker.id;
            const std::size_t moveIndex = moves.size();
            moves.push_back(move);
            leaveMarkerByStep.emplace(fromKey, departMarker.id);
            enterMarkerByStep.emplace(toKey, arriveMarker.id);
            leaveMarkerToMoveIndex.emplace(departMarker.id, moveIndex);

            const double safety = clampSafetyDistance(config.safetyDistance, distance);
            const double midDistance = std::max(distance - 2.0 * safety, 0.0);
            const double lowerBoundDepart = (safety > 0.0) ? (safety / speed) : 0.0;
            const double lowerBoundMid = (midDistance > 0.0) ? (midDistance / speed) : 0.0;
            const double lowerBoundArrive = (safety > 0.0) ? (safety / speed) : 0.0;

            TemporalConstraint departConstraint;
            departConstraint.fromEventId = fromEventId;
            departConstraint.toEventId = departMarker.id;
            departConstraint.type = TemporalConstraintType::Safety;
            departConstraint.lowerBound = lowerBoundDepart;
            constraints.push_back(departConstraint);

            TemporalConstraint travelConstraint;
            travelConstraint.fromEventId = departMarker.id;
            travelConstraint.toEventId = arriveMarker.id;
            travelConstraint.type = TemporalConstraintType::Safety;
            travelConstraint.lowerBound = lowerBoundMid;
            constraints.push_back(travelConstraint);

            TemporalConstraint arriveConstraint;
            arriveConstraint.fromEventId = arriveMarker.id;
            arriveConstraint.toEventId = toEventId;
            arriveConstraint.type = TemporalConstraintType::Safety;
            arriveConstraint.lowerBound = lowerBoundArrive;
            constraints.push_back(arriveConstraint);
        }
    }

    for (auto& relation : waitRelations) {
        auto leaveIt = leaveMarkerByStep.find(relation.awaited);
        auto enterIt = enterMarkerByStep.find(relation.waiter);

        TemporalConstraint& constraint = constraints[relation.constraintIndex];
        constraint.fromEventId = relation.awaitedEventId;
        constraint.toEventId = relation.waiterEventId;

        if (leaveIt != leaveMarkerByStep.end()) {
            constraint.fromEventId = leaveIt->second;
        }
        if (enterIt != enterMarkerByStep.end()) {
            constraint.toEventId = enterIt->second;
        }
    }

    QSet<QString> waitSeen;
    for (const auto& relation : waitRelations) {
        WaitConstraint wait;
        wait.awaitedAgv = relation.awaited.agvNumber;
        wait.waiterAgv = relation.waiter.agvNumber;
        wait.waiterStationIndex = std::max(0, relation.waiter.routeIndex - 1);

        QString unlockFrom;
        QString unlockTo;
        auto leaveIt = leaveMarkerByStep.find(relation.awaited);
        if (leaveIt != leaveMarkerByStep.end()) {
            auto moveIdxIt = leaveMarkerToMoveIndex.find(leaveIt->second);
            if (moveIdxIt != leaveMarkerToMoveIndex.end()) {
                const MoveInfo& move = moves[moveIdxIt->second];
                unlockFrom = move.fromStation;
                unlockTo = move.toStation;
            }
        }
        if (unlockFrom.isEmpty() || unlockTo.isEmpty() || unlockFrom == unlockTo) {
            unlockFrom = relation.station;
            unlockTo = relation.station;
        }

        if (unlockFrom == unlockTo) {
            continue;
        }

        wait.unlockFrom = unlockFrom;
        wait.unlockTo = unlockTo;

        const QString key = QString::number(wait.awaitedAgv) + "|" +
                            QString::number(wait.waiterAgv) + "|" + wait.unlockFrom + "|" +
                            wait.unlockTo + "|" + QString::number(wait.waiterStationIndex);
        if (waitSeen.contains(key)) {
            continue;
        }
        waitSeen.insert(key);
        result.waitConstraints.append(wait);
        result.specialStationWhiteList.insert(relation.station, 1);
    }

    const int eventCount = static_cast<int>(events.size());
    std::vector<double> earliest(eventCount, 0.0);
    std::vector<int> indegree(eventCount, 0);
    std::vector<std::vector<std::pair<int, double>>> adjacency(eventCount);

    for (const auto& constraint : constraints) {
        if (constraint.fromEventId < 0 || constraint.toEventId < 0) {
            continue;
        }
        const int from = constraint.fromEventId;
        const int to = constraint.toEventId;
        if (from < 0 || from >= eventCount || to < 0 || to >= eventCount) {
            continue;
        }
        adjacency[from].push_back({to, std::max(0.0, constraint.lowerBound)});
        indegree[to] += 1;
    }

    std::queue<int> queue;
    for (int evt = 0; evt < eventCount; ++evt) {
        if (indegree[evt] == 0) {
            queue.push(evt);
        }
    }

    int processed = 0;
    while (!queue.empty()) {
        const int current = queue.front();
        queue.pop();
        ++processed;
        for (const auto& edge : adjacency[current]) {
            const int next = edge.first;
            const double lb = edge.second;
            const double candidate = earliest[current] + lb;
            if (candidate > earliest[next]) {
                earliest[next] = candidate;
            }
            indegree[next] -= 1;
            if (indegree[next] == 0) {
                queue.push(next);
            }
        }
    }

    if (processed != eventCount) {
        // Graph contains a cycle; fall back to zero schedule.
        std::fill(earliest.begin(), earliest.end(), 0.0);
    }

    for (const auto& evt : events) {
        result.schedule.eventTimes.insert(evt.id, earliest[evt.id]);
    }

    result.events.clear();
    result.events.reserve(static_cast<int>(events.size()));
    for (const auto& evt : events) {
        result.events.append(evt);
    }

    result.constraints.clear();
    result.constraints.reserve(static_cast<int>(constraints.size()));
    for (const auto& c : constraints) {
        result.constraints.append(c);
    }

    return result;
}

}  // namespace pathfinding
}  // namespace application
